""" Calculate Net Benefit
"""

import os
from io import open
import tempfile
import shutil
import numpy as np
import copy as cp
import pandas as pd
import itertools


# SETTINGS  =========================================================================================================================

FNAME_IN  = 'data/sub_info.csv'
FNAME_OUT = 'results/04-calculated_net_benefits.csv'
FNAME_OUT1 = 'data/calculated_net_benefits.csv' #output to the intermediate folder for use by R 05_ file

# MAIN FUNCTION  ===============================================================================================================================================

def main():
    s1 = np.arange(.0, 1.01, step=.025)
    iterables = [s1, s1, s1, s1]
    DISTRIBUTIONS = []
    for t in itertools.product(*iterables):
        DISTRIBUTIONS.append(np.asarray(t))
    DISTRIBUTIONS = np.array(DISTRIBUTIONS)
    DISTRIBUTIONS = DISTRIBUTIONS[np.round(DISTRIBUTIONS.sum(axis=1), 3) == 1,]

    all_data = pd.read_csv(FNAME_IN)
    all_data.head()

    # Get number of all subjects
    n_subs = all_data.shape[0]

    out_data = []

    actionsByPeriod = np.ones((501, len(all_data.eid))) * (-1)
    subNames = []

    # For each subject eid go to corresponding game data .csv file and import the data into output variables
    for i, eid in enumerate(all_data.eid):

        sub_data = all_data[all_data.eid == eid]
        print("Subject# ", i, " eid=", eid)

        subNames.append(eid)

        game_data = pd.read_csv('experiments/' + str(eid) + '.csv', header=None)
        game_data = game_data.ix[1:]
        game_data.columns = ['eid', 'config', 'round_no', 'moves_taken', '1st_waitimes',
                             'last_action', 'balance', 'position0', 'position1',
                             'prize_pos0', 'prize_pos1', 'atPrize', 'atStart', 'atDoor']
        game_data['atDoor1'] = (game_data.atDoor == 1) & (game_data.position0 == 4)
        game_data['atDoor2'] = ((game_data.atDoor == 1) & (game_data.position0 == 2)) | (
                    (game_data.atDoor == 1) & (game_data.position0 == 0))
        game_data['Break1'] = np.min(game_data.moves_taken[(game_data.position0 == 5) & (game_data.position1 == 3)])

        temp = np.array(game_data['last_action'])
        temp = np.append(temp, [-99])
        game_data['next_action'] = pd.Series(temp)
        game_data['moves_remaining'] = 500 - game_data['moves_taken']

        temp = np.array(game_data['atPrize'])[:-1]
        temp = np.append([-99, -99], temp)
        game_data['prevPrize'] = pd.Series(temp)

        prizeLoc = 0 * ((game_data.prize_pos0 == 6) & (game_data.prize_pos1 == 3)) + 1 * (
                    (game_data.prize_pos0 == 3) & (game_data.prize_pos1 == 6)) + \
                   2 * ((game_data.prize_pos0 == 0) & (game_data.prize_pos1 == 3)) + 3 * (
                               (game_data.prize_pos0 == 3) & (game_data.prize_pos1 == 0))
        temp = np.array(prizeLoc)[:-1]
        temp = np.append([-99, -99], temp)
        game_data['prevPrizeLoc'] = pd.Series(temp)

        benefitData = game_data[game_data.atStart == 1]
        benefitData['netBenefit'] = 0
        benefitData['Door2HitsUntilNow'] = 0
        benefitData['prizeSequence'] = ""

        benefitData['Incentives'] = int((np.min(sub_data.Config) - 1) % 3)
        benefitData['Structures'] = int(np.floor((sub_data.Config - 1) / 3))

        PRIOR = np.repeat(1.0 / DISTRIBUTIONS.shape[0], DISTRIBUTIONS.shape[0])
        PROBS = np.dot(PRIOR, DISTRIBUTIONS)
        actSeq = ""
        for index, row in benefitData.iterrows():

            if row['prevPrize'] == 1:
                # CALCULATE PROBS
                d = row['prevPrizeLoc']
                actSeq = actSeq + 'trbl'[d]
                POSTERIOR = DISTRIBUTIONS[:, d] * PRIOR / UNCONDITIONAL[d]
                PRIOR = POSTERIOR

            if row['atStart'] == 1:
                actionsByPeriod[row['moves_taken'], len(subNames) - 1] = row['next_action']

            UNCONDITIONAL = np.dot(PRIOR, DISTRIBUTIONS)

            TOP_RIGHT = np.dot([3, 9, 15, 21], UNCONDITIONAL.T)
            TOP_LEFT = np.dot([3, 21, 15, 9], UNCONDITIONAL.T)
            BOT_RIGHT = np.dot([15, 9, 3, 21], UNCONDITIONAL.T)
            BOT_LEFT = np.dot([15, 21, 3, 9], UNCONDITIONAL.T)

            TOP = np.minimum(TOP_RIGHT, TOP_LEFT)
            BOT = np.minimum(BOT_RIGHT, BOT_LEFT)

            HORIZON_TOP = np.floor(row['moves_remaining'] / TOP)
            HORIZON_BOT = np.floor(row['moves_remaining'] / BOT)

            benefitData.netBenefit[index] = (HORIZON_BOT - HORIZON_TOP) * 6
            benefitData.Door2HitsUntilNow[index] = np.sum(
                (game_data.atDoor == 1) & (game_data.moves_taken > game_data.Break1) & (
                            game_data.moves_taken <= row['moves_taken']))
            benefitData.prizeSequence[index] = actSeq

        if len(out_data) == 0:
            out_data = cp.deepcopy(benefitData)
        else:
            out_data = out_data.append(benefitData, ignore_index=True)

    out_data.head()

    out_data.to_csv(FNAME_OUT, index=False, header=True)
    out_data.to_csv(FNAME_OUT1, index=False, header=True)


# COMMANDLINE  ===============================================================================================================================================
if __name__ == "__main__": main()